#include <stdint.h>

#include <inject.h>

#define JUMP_SIZE (6 + sizeof(void*))

#define TRAMPOLINE_SIZE (6 + sizeof(void*))
#define ENTRY_SIZE 5

// Original values to recover after the injection
// Recovery is performed by the assembly payload
#pragma pack(push, 1)
struct recovery_data {
    void *trampoline_address;
    char trampoline_data[TRAMPOLINE_SIZE];

    void *entry_point_address;
    char entry_point_data[ENTRY_SIZE];

    void *import_descriptor_address;
    IMAGE_IMPORT_DESCRIPTOR import_descriptor_data;

    void *size_field_address;
    DWORD size_field_data;
};
#pragma pack(pop)

static inline void write_protected_process_memory(HANDLE process, void *address, const void *buf, size_t size) {
    DWORD old_protect;
    VirtualProtectEx(process, address, size, PAGE_EXECUTE_READWRITE, &old_protect);

    WriteProcessMemory(process, address, buf, size, NULL);

    VirtualProtectEx(process, address, size, old_protect, &old_protect);
}

void inject(HANDLE process, const void *payload, size_t payload_size, const wchar_t *dllPath) {
    // Find the EXE header in the process
    char exe_header[1024];
    IMAGE_DOS_HEADER *dos_header = NULL;
    IMAGE_NT_HEADERS64 *nt_headers = NULL;

    MEMORY_BASIC_INFORMATION memory_info;
    char *current_address = 0x0;
    while (VirtualQueryEx(process, current_address, &memory_info, sizeof(memory_info))) {
        ReadProcessMemory(process, current_address, exe_header, sizeof(exe_header), NULL);

        dos_header = (IMAGE_DOS_HEADER*)exe_header;

        // DOS header magic "MZ"
        if (dos_header->e_magic != 0x5A4D) {
            goto cont;
        } 

        nt_headers = (IMAGE_NT_HEADERS64*)(exe_header + dos_header->e_lfanew);

        // NT header signature "PE"
        if (nt_headers->Signature != 0x4550) {
            goto cont;
        }

        // Skip DLLs
        if ((nt_headers->FileHeader.Characteristics & IMAGE_FILE_DLL) == IMAGE_FILE_DLL) {
            goto cont;
        }

        // Skip potential headers without an entry point 
        // I have no idea how and why they exist, but apparently they do
        if (nt_headers->OptionalHeader.AddressOfEntryPoint == 0) {
            goto cont;
        }

        // Found EXE header
        break;

    cont:
        current_address += memory_info.RegionSize;
    }

    char *exe = (char*)memory_info.BaseAddress;


    // Inject the loader into the process
    size_t dll_path_size = (wcslen(dllPath) + 1) * sizeof(wchar_t);

    size_t alloc_size = payload_size + dll_path_size + sizeof(struct recovery_data);
    char *remote_alloc = VirtualAllocEx(process, NULL, alloc_size, MEM_COMMIT, PAGE_EXECUTE_READWRITE);

    // Write the assembly payload and dll path
    WriteProcessMemory(process, remote_alloc, payload, payload_size, NULL);
    WriteProcessMemory(process, remote_alloc + payload_size, dllPath, dll_path_size, NULL);


    // Modify the executable to run the assembly payload
    // Recovery data structure
    struct recovery_data rd;

    // Find a place for the trampoline to the assembly payload
    char *trampoline = NULL;
    IMAGE_SECTION_HEADER *sections = (IMAGE_SECTION_HEADER*)(nt_headers + 1);
    for (size_t i = 0; i < nt_headers->FileHeader.NumberOfSections; i++) {
        if (sections[i].Characteristics & IMAGE_SCN_MEM_EXECUTE) {
            trampoline = exe + sections[i].VirtualAddress;
        }
    }

    // Save original trampoline bytes
    rd.trampoline_address = trampoline;
    ReadProcessMemory(process, rd.trampoline_address, rd.trampoline_data, sizeof(rd.trampoline_data), NULL);

    // Write the trampoline
    char trampoline_payload[TRAMPOLINE_SIZE] = { 0xFF, 0x25, 0x00, 0x00, 0x00, 0x00 };
    memcpy(trampoline_payload + 6, &remote_alloc, sizeof(remote_alloc));

    write_protected_process_memory(process, trampoline, trampoline_payload, sizeof(trampoline_payload));


    // Save original entry point bytes
    char *entry_point = exe + nt_headers->OptionalHeader.AddressOfEntryPoint;

    rd.entry_point_address = entry_point;
    ReadProcessMemory(process, rd.entry_point_address, rd.entry_point_data, sizeof(rd.entry_point_data), NULL);

    // Replace the entry point with a jump to the trampoline
    int32_t trampoline_offset = (int32_t)((ptrdiff_t)trampoline - (ptrdiff_t)entry_point - 5);

    char entry_payload[ENTRY_SIZE] = { 0xE9 };
    memcpy(entry_payload + 1, &trampoline_offset, sizeof(trampoline_offset));
    
    write_protected_process_memory(process, entry_point, entry_payload, sizeof(entry_payload));


    // Break the import table to prevent any dlls from being loaded
    // Step 1: break the first import descriptor
    char *import_descriptors = exe + nt_headers->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress;

    // Save the original descriptor address and bytes
    rd.import_descriptor_address = import_descriptors;
    ReadProcessMemory(process, rd.import_descriptor_address, &rd.import_descriptor_data, sizeof(rd.import_descriptor_data), NULL);

    // Overwrite with zeroes    
    IMAGE_IMPORT_DESCRIPTOR first_descriptor = { 0 };
    write_protected_process_memory(process, import_descriptors, &first_descriptor, sizeof(first_descriptor));

    // Step 2: break the image data directory entry
    char* dd_addr = ((char*)&(nt_headers->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].Size)) - exe_header + exe;
    
    // Save the original value
    rd.size_field_address = dd_addr;
    ReadProcessMemory(process, rd.size_field_address, &rd.size_field_data, sizeof(rd.size_field_data), NULL);

    // Set to 0
    DWORD new_size = 0;
    write_protected_process_memory(process, dd_addr, &new_size, sizeof(new_size));

    // Write recovery data to the allocation
    WriteProcessMemory(process, remote_alloc + payload_size + dll_path_size, &rd, sizeof(rd), NULL);
}
